from .base import MetricBase
import torch
import numpy as np
from scipy.optimize import linear_sum_assignment

class OSPA(MetricBase):
    def __init__(self, cutoff=10, p=1):
        super().__init__()
        self.cutoff = cutoff
        self.p = p
        assert self.p > 0 and self.cutoff > 0
    
    def compute(self, y_true, y_pred):
        batch_size = y_true.size(0)
        ospa_dist = []
        for b in range(batch_size):
            true_positions = torch.nonzero(y_true[b].squeeze()).cpu().numpy()
            pred_positions = torch.nonzero(y_pred[b].squeeze()).cpu().numpy()
            
            m, n = len(true_positions), len(pred_positions)
            
            if m == 0 and n == 0:
                ospa_dist.append(0)
                continue
            
            if m == 0 or n == 0:
                ospa_dist.append((self.cutoff ** self.p) ** ( 1 / self.p))
                continue
            
            dist_matrix = torch.cdist(
                torch.tensor(true_positions, dtype=torch.float32),
                torch.tensor(pred_positions, dtype=torch.float32)
            ).numpy()
            
            dist_matrix = np.minimum(dist_matrix, self.cutoff)
            row_ind, col_ind = linear_sum_assignment(dist_matrix)
            matching_cost = dist_matrix[row_ind, col_ind].sum()
            unmatched_cost = (self.cutoff ** self.p) * abs(m - n)
            ospa_distance = ((matching_cost + unmatched_cost) / max(m, n)) ** (1 / self.p)
            ospa_dist.append(ospa_distance)
        ospa = sum(ospa_dist) / len(ospa_dist)
        return ospa
    
    def summary(self):
        ospa = self.value / self.total
        print(f'OSPA: {ospa}')
        return ospa.item()
    
if __name__ == '__main__':
    ospa = OSPA()
    y_true = torch.stack([torch.bernoulli(torch.full((64,64), 0.3)).unsqueeze(0) for _ in range(16)]).float()
    y_pred = torch.randint(0,2, size=(16,1,64,64)).float()
    print(y_true.shape)
    ospa(y_true, y_true)
    ospa.summary()